import tqdm
from gxl_ai_utils.utils import utils_file

input_shard_path ="/home/work_nfs14/asr_data/data/wenetspeech/shards/wenetspeech_new_all"
input_text_path = "/home/work_nfs14/asr_data/data/wenetspeech/shards/wenetspeech_new_all/text"
# 得到 train_shards.list 和 dev_shards.list
def do_fuc1():
    tar_file_list = utils_file.get_list_for_wav_dir(input_shard_path, suffix='tar')
    output_train_path = './train_shards.list'
    output_dev_path = './dev_shards.list'
    train_list = tar_file_list[:-10]
    dev_list = tar_file_list[-10:]
    utils_file.write_list_to_file(train_list, output_train_path)
    utils_file.write_list_to_file(dev_list, output_dev_path)


# 得到词表
def do_sort_dict_by_value(my_dict, asc=True):
    """
    对字典按照value排序,默认升序, 升序是将小的数排在前面
    """
    if asc:
        return dict(sorted(my_dict.items(), key=lambda x: x[1]))
    else:
        return dict(sorted(my_dict.items(), key=lambda x: x[1], reverse=True))
def do_get_vocab_dict_from_text_scp_by_char_fn(text_scp_path, output_vocab_path=None):
    """
    从一个text.scp中得到一个token.txt, 按照wenet u2++需要的格式
    """
    token_dict = {}
    text_dict = utils_file.load_dict_from_scp(text_scp_path)
    for key, text in tqdm.tqdm(text_dict.items(), desc='get vocab', total=len(text_dict)):
        for char in text:
            if char in [' ','\n','\t']:
                continue
            if char not in token_dict:
                token_dict[char] = 0
            else:
                token_dict[char] += 1
    sort_dict = do_sort_dict_by_value(token_dict, False)
    res_dict = {}
    res_dict['<blank>'] = 0
    res_dict['<unk>'] = 1
    res_dict['<sos/eos>'] = 2
    for key, value in sort_dict.items():
        res_dict[key] = len(res_dict)
    if output_vocab_path:
        utils_file.write_dict_to_scp(res_dict, output_vocab_path)
    else:
        return res_dict

# do_get_vocab_dict_from_text_scp_by_char_fn(input_text_path, './token.txt')
def do_test():
    res_dict = utils_file.load_dict_from_scp('./token.txt')
    res_dict = utils_file.get_subdict(res_dict,2500,2700)
    utils_file.print_dict(res_dict)

def do_musan():
    path1 = "/home/work_nfs8/asr_data/data/musan/music/wav.scp"
    path2 = "/home/work_nfs8/asr_data/data/musan/noise/wav.scp"
    wav_dict = utils_file.load_dict_from_scp(path1)
    new_wav_dict = {}
    for key, value in wav_dict.items():
        new_wav_dict[key] = value + ".wav"
    utils_file.write_dict_to_scp(new_wav_dict, path1)

    wav_dict = utils_file.load_dict_from_scp(path2)
    new_wav_dict = {}
    for key, value in wav_dict.items():
        new_wav_dict[key] = value + ".wav"
    utils_file.write_dict_to_scp(new_wav_dict, path2)




if __name__ == '__main__':
    do_musan()



